# -*- coding: utf-8 -*-
import argparse
import os
from os.path import join

parser = argparse.ArgumentParser(description='Forward all testing images.')
parser.add_argument('--method',
                    type=str,
                    default='poolnet_pro-res18_reweight2')
parser.add_argument('--epoch', type=int, default=20)
parser.add_argument('--traverse', type=int, default=1)
parser.add_argument('--range', type=str, default='100,101')
parser.add_argument('--gt_dir', type=str, default='./datasets/sal')
parser.add_argument('--save_dir', type=str, default='./results')
# parser.add_argument('--sal_dir', type=str, default='CSNet/results/')
args = parser.parse_args()

modelprefix = args.method

outputs = ['ECSSD/test',]#,'ECSSD/val',]
# LOG_FOUT = open(
#     join(args.save_dir, modelprefix,
#          "FmeasureResults_" + str(args.method) + '_' + str(outputs) + '.txt'),
#     'a')


# def log_string(out_str):
#     out_str = str(out_str)
#     LOG_FOUT.write(out_str + '\n')
#     LOG_FOUT.flush()
#     print(out_str)



# #########eval########### #
if args.traverse:
    testrange = args.range.split(',')
    startepoch = int(testrange[0])
    endepoch = int(testrange[1])
else:
    startepoch = args.epoch
    endepoch = args.epoch + 1
bestF = 0
best_epoch = 0
print("Start:", startepoch, "End:", endepoch)
for thisepoch in range(startepoch, endepoch):
    for output in outputs:
        save_dir = join(args.save_dir, modelprefix,
                        output.split('/')[0] + '_' + str(thisepoch))
        if not os.path.isdir(save_dir):
            print(save_dir + " not exist.")
            continue
        else:
            print(output + " Eval epoch: " + str(thisepoch))
        gt = join(args.gt_dir, output, 'GT')
        # sal = join(args.sal_dir, modelprefix, output + '_' + str(thisepoch))
        # res = os.listdir(save_dir)
        # vallist = join(args.save_dir, modelprefix,
        #                "val_" + output + '_' + str(thisepoch) + '.txt')
        # fobj = open(vallist, 'w')
        # for sub in res:
        #     fobj.write(join(save_dir, sub) + ' ' + join(gt, sub) + '\n')
        # fobj.close()
        
        valreslist = join(
            args.save_dir, modelprefix,
            "FmeasureResult_" + output.split('/')[0] + '_' + str(thisepoch) + ".txt")
        valresobj = open(valreslist, 'w')
        # print(vallist+" generate done.")
        valcmd =  "python ./salmetric.py "+'--sal_dir' + ' ' + save_dir + ' --gt_dir' + ' ' + gt
        print(valcmd)
        content = os.popen(valcmd).read()
        valresobj.write(content)
        results = content.split('\n')[-8:]
        print(results)
        thisMaxF = float(results[0].split()[1])
        if bestF < thisMaxF:
            bestF = thisMaxF
            best_epoch = thisepoch
        print(valreslist + " eval done.")
        valresobj.close()
print("BestF: " + str(bestF) + " in Epoch: " + str(best_epoch))
